/*
 * CLMP.C
 *
 * Multi processor support for older versions of Visual C++ that don't implement
 * it natively.
 *
 * Copyright (c) 2015 Malcolm J. Smith
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#define WIN32_LEAN_AND_MEAN 1
#include <windows.h>
#include <stdlib.h>
#include <stdio.h>
#include <tchar.h>

#ifdef MINICRT
#include <minicrt.h>
#endif

#if (defined(_MSC_VER) && _MSC_VER >= 1200 && _MSC_VER < 1310)
#pragma comment(linker, "/OPT:NOWIN98")
#endif

DWORD GlobalExitCode = 0;

typedef struct _CLMP_PIPE_BUFFER {
    HANDLE Pipe;
    PUCHAR ReadBuffer;
    DWORD  ReadBufferSize;
} CLMP_PIPE_BUFFER, *PCLMP_PIPE_BUFFER;

typedef struct _CLMP_PROCESS_INFO {
    PROCESS_INFORMATION WindowsProcessInfo;
    CLMP_PIPE_BUFFER Pipes[2];
    LPTSTR Filename;
} CLMP_PROCESS_INFO, *PCLMP_PROCESS_INFO;

VOID
ProcessSinglePipeOutput (
    PCLMP_PROCESS_INFO Process,
    BOOL Final
    )
{
    DWORD BytesToRead = 0;
    DWORD BytesToAllocate;
    DWORD PipeNum;
    PCLMP_PIPE_BUFFER Pipe;
    HANDLE OutHandle;

    //
    //  Each outstanding process has two handles to read from: one
    //  corresponding to its standard output, the other to standard
    //  error.  We read from each and write to the corresponding
    //  handle from this process.
    //
    //  If this seems like a lot of work, the reason is because we
    //  want to write line-by-line to ensure we don't get garbage
    //  output caused by racing child processes.  It also allows us
    //  to give the user more information about what's failing.
    //

    for (PipeNum = 0; PipeNum < sizeof(Process->Pipes)/sizeof(Process->Pipes[0]); PipeNum++) {

        Pipe = &Process->Pipes[PipeNum];
        if (PipeNum == 0) {
            OutHandle = GetStdHandle(STD_OUTPUT_HANDLE);
        } else {
            OutHandle = GetStdHandle(STD_ERROR_HANDLE);
        }

        //
        //  Check if we have an initialized process and if we have data to read.
        //

        if (Pipe->Pipe != NULL) {

            BytesToRead = 0;
            PeekNamedPipe( Pipe->Pipe, NULL, 0, NULL, &BytesToRead, NULL );

            if (BytesToRead != 0) {

                PUCHAR NewBuffer;
                DWORD  BytesRead;
                DWORD  ThisChar;

                //
                //  We allocate a buffer for any leftover data (that didn't have a \n)
                //  plus whatever new data we have now.  Copy forward the old data and
                //  read the new data.
                //

                BytesToAllocate = Pipe->ReadBufferSize + BytesToRead;

                NewBuffer = HeapAlloc(GetProcessHeap(), 0, BytesToAllocate);

                if (Pipe->ReadBufferSize) {
                    memcpy(NewBuffer, Pipe->ReadBuffer, Pipe->ReadBufferSize);
                }

                ReadFile(Pipe->Pipe, NewBuffer + Pipe->ReadBufferSize, BytesToRead, &BytesRead, NULL);

                //
                //  Go look for \n's now.  If we find one, send that data to the correct
                //  handle for this process, and continue searching for more.
                //

                for (ThisChar = 0; ThisChar < BytesToAllocate; ThisChar++) {

                    if (NewBuffer[ThisChar] == '\n') {

                        WriteFile(OutHandle, NewBuffer, ThisChar + 1, &BytesRead, NULL);
                        memmove(NewBuffer, NewBuffer + ThisChar + 1, BytesToAllocate - ThisChar - 1);
                        BytesToAllocate -= ThisChar + 1;
                        ThisChar = 0;
                    }
                }

                //
                //  If we had any previous leftover data, it's moot now.  Tear down
                //  that buffer.
                //

                if (Pipe->ReadBuffer != NULL) {
                    HeapFree(GetProcessHeap(), 0, Pipe->ReadBuffer);
                    Pipe->ReadBuffer = NULL;
                    Pipe->ReadBufferSize = 0;
                }

                //
                //  If we still have data, either carry it over or if the process is
                //  dead just write what we have.  If we don't have data, tear down
                //  our buffer.
                //

                if (BytesToAllocate > 0) {
                    if (Final) {
                        WriteFile(OutHandle, NewBuffer, BytesToAllocate, &BytesRead, NULL);
                        HeapFree(GetProcessHeap(), 0, NewBuffer);
                    } else {
                        Pipe->ReadBuffer = NewBuffer;
                        Pipe->ReadBufferSize = BytesToAllocate;
                    }
                } else {
                    HeapFree(GetProcessHeap(), 0, NewBuffer);
                }
            }
        }
    }
}

VOID
ProcessPipeOutput (
    PCLMP_PROCESS_INFO Processes,
    DWORD NumberProcesses
    )
{
    DWORD i;

    //
    //  Look for outstanding data in any process and write it out.
    //  

    for (i = 0; i < NumberProcesses; i++) {
        ProcessSinglePipeOutput( &Processes[i], FALSE );
    }
}

//
//  Wait for a child process.  If it failed, we fail.
//

VOID
WaitOnProcess (
    PCLMP_PROCESS_INFO Process
    )
{
    DWORD ExitCode;
    DWORD PipeNum;

    while (Process->WindowsProcessInfo.hProcess != NULL) {

        HANDLE WaitHandles[3];
        DWORD  HandleNum;

        //
        //  We want to wait for the process to tear down, but we also need
        //  to break out if the process is generating output that we need
        //  to handle, lest we might prevent it from tearing down.
        //

        WaitHandles[0] = Process->WindowsProcessInfo.hProcess;
        WaitHandles[1] = Process->Pipes[0].Pipe;
        WaitHandles[2] = Process->Pipes[1].Pipe;

        if (WaitHandles[2] == NULL && WaitHandles[1] == NULL) {
            HandleNum = 1;
        } else {
            HandleNum = 3;
        }

        HandleNum = WaitForMultipleObjects(HandleNum, WaitHandles, FALSE, INFINITE);
        if (HandleNum > WAIT_OBJECT_0 && HandleNum <= WAIT_OBJECT_0 + 2) {
            ProcessSinglePipeOutput(Process, FALSE);
        } else {
            ProcessSinglePipeOutput(Process, TRUE);
            GetExitCodeProcess( Process->WindowsProcessInfo.hProcess, &ExitCode );

            //
            //  If a child failed and the parent is still going, fail with
            //  the same code.
            //

            if (ExitCode != 0 && GlobalExitCode == 0) {
                GlobalExitCode = ExitCode;
            }

            //
            //  Clean up and tear down so the process slot can be reused as
            //  necessary.
            //

            CloseHandle(Process->WindowsProcessInfo.hProcess);
            CloseHandle(Process->WindowsProcessInfo.hThread);
        
            Process->WindowsProcessInfo.hProcess = NULL;
            Process->WindowsProcessInfo.hThread = NULL;

            for (PipeNum = 0; PipeNum < sizeof(Process->Pipes)/sizeof(Process->Pipes[0]); PipeNum++) {

                CloseHandle(Process->Pipes[PipeNum].Pipe);

                Process->Pipes[PipeNum].Pipe = NULL;

                if (Process->Pipes[PipeNum].ReadBuffer) {
                    HeapFree(GetProcessHeap(), 0, Process->Pipes[PipeNum].ReadBuffer);
                    Process->Pipes[PipeNum].ReadBuffer = NULL;
                }
                Process->Pipes[PipeNum].ReadBufferSize = 0;
            }
        }
    }
}


#ifdef UNICODE
#define ENTRYPOINT wmain
#else
#define ENTRYPOINT main
#endif

int __cdecl
ENTRYPOINT (
    int argc,
    LPTSTR argv[]
    )
{
    TCHAR szCmdCommon[1024];
    TCHAR szCmdComplete[1024];
    int i, j;
    PCLMP_PROCESS_INFO ProcessInfo;
    DWORD CurrentProcess = 0;
    DWORD NumberProcesses = 0;
    SYSTEM_INFO SysInfo;
    BOOLEAN MultiProcPossible = FALSE;
    BOOLEAN MultiProcNotPossible = FALSE;

    GetSystemInfo(&SysInfo);

    _tcscpy(szCmdCommon, _T("cl "));

    //
    //  Scan the command line looking for switches that should
    //  be common to all of the processes we spawn.  Also look
    //  for switches which tell us that multi processing is
    //  incompatible with the requested operation, and in that
    //  case, disable it.
    //

    for (i=1;i<argc;i++) {
        if (argv[i][0]=='-' || argv[i][0]=='/') {

            //
            //  Check for MP switch and adjust the number of processes
            //  as requested.  Don't tell the compiler about this.
            //

            if (argv[i][1]=='M' &&
                argv[i][2]=='P') {

                if (argv[i][3]>='0' && argv[i][3]<='9') {

                    NumberProcesses = _ttoi(&argv[i][3]);
                }

            } else {
                _tcscat(szCmdCommon, _T(" "));
                _tcscat(szCmdCommon, argv[i]);
            }

            //
            //  Compile without linking - we need this for multiproc, because
            //  if we're linking in the same pass issuing different processes
            //  with different portions would generate the wrong result
            //

            if (argv[i][1]=='c') {
                MultiProcPossible = TRUE;
            }

            //
            //  Preprocess to stdout - we can't do this in parallel without
            //  generating garbage
            //

            if (argv[i][1]=='E') {
                MultiProcNotPossible = TRUE;
            }

            //
            //  Generating debug into a PDB - we can't do this in parallel
            //  or we'll get sharing violations on the PDB
            //

            if (argv[i][1]=='Z') {
                for (j = 2; argv[i][j]!='\0'; j++) {
                    if (argv[i][j] == 'i' ||
                        argv[i][j] == 'I') {

                        MultiProcNotPossible = TRUE;
                    }
                }
            }

            //
            //  Precompiled header - we can't do this in parallel or we'll
            //  get sharing violations on the precompiled header file
            //

            if (argv[i][1]=='Y') {
                for (j = 2; argv[i][j]!='\0'; j++) {
                    if (argv[i][j] == 'c' ||
                        argv[i][j] == 'X') {

                        MultiProcNotPossible = TRUE;
                    }
                }
            }
        }
    }

    //
    //  If we disabled multi processing, we still want to have one child
    //  or we won't get far.
    //

    if (!MultiProcPossible || MultiProcNotPossible) {
        NumberProcesses = 1;
    } else if (NumberProcesses == 0) {
        NumberProcesses = SysInfo.dwNumberOfProcessors + 1;
    }

    ProcessInfo = HeapAlloc(GetProcessHeap(), 0, sizeof(CLMP_PROCESS_INFO) * NumberProcesses);

    if (ProcessInfo == NULL) {
        return EXIT_FAILURE;
    }

    ZeroMemory(ProcessInfo, sizeof(CLMP_PROCESS_INFO) * NumberProcesses);

    //
    //  Scan again looking for source files, and spawn one child
    //  process per argument found, combined with the command
    //  line we built earlier.
    //

    for (i=1;i<argc;i++) {
        if (argv[i][0]!='-' && argv[i][0]!='/') {

            STARTUPINFO StartupInfo;
            DWORD MyProcess = CurrentProcess;
            SECURITY_ATTRIBUTES SecurityAttributes;
            HANDLE WriteOutPipe, WriteErrPipe;

            _stprintf(szCmdComplete, _T("%s %s"), szCmdCommon, argv[i]);
            ZeroMemory( &StartupInfo, sizeof(StartupInfo) );
            StartupInfo.cb = sizeof(StartupInfo);

            MyProcess %= NumberProcesses;

            //
            //  Look for pending output in any process.  This is a little
            //  gratuitous, but a child may block waiting for us if we don't,
            //  so do it now.
            //

            ProcessPipeOutput(ProcessInfo, NumberProcesses);

            //
            //  If we've run out of processors, wait for the next child
            //  process and reuse its slot.
            //

            if (CurrentProcess >= NumberProcesses) {
                WaitOnProcess( &ProcessInfo[MyProcess] );
                if (GlobalExitCode) {
                    goto drain;
                }
            }

            //
            //  We need to specify security attributes because we want our
            //  standard output and standard error handles to be inherited.
            //

            ZeroMemory( &SecurityAttributes, sizeof(SecurityAttributes) );

            SecurityAttributes.nLength = sizeof(SecurityAttributes);
            SecurityAttributes.bInheritHandle = TRUE;

            //
            //  Create the aforementioned handles.
            //

            if (!CreatePipe( &ProcessInfo[MyProcess].Pipes[0].Pipe, &WriteOutPipe, &SecurityAttributes, 0)) {

                GlobalExitCode = EXIT_FAILURE;
                goto drain;
            }

            if (!CreatePipe( &ProcessInfo[MyProcess].Pipes[1].Pipe, &WriteErrPipe, &SecurityAttributes, 0)) {

                GlobalExitCode = EXIT_FAILURE;
                goto drain;
            }

            ProcessInfo[MyProcess].Filename = argv[i];

            //
            //  The child process should write to the write handles, but
            //  this process doesn't want those, so we close them immediately
            //  below.
            //

            StartupInfo.dwFlags = STARTF_USESTDHANDLES;
            if (GetStdHandle(STD_OUTPUT_HANDLE) == GetStdHandle(STD_ERROR_HANDLE)) {
                StartupInfo.hStdOutput = WriteOutPipe;
                StartupInfo.hStdError = WriteOutPipe;
            } else {
                StartupInfo.hStdOutput = WriteOutPipe;
                StartupInfo.hStdError = WriteErrPipe;
            }

            if (!CreateProcess( NULL, szCmdComplete, NULL, NULL, TRUE, 0, NULL, NULL, &StartupInfo, &ProcessInfo[MyProcess].WindowsProcessInfo )) {
                GlobalExitCode = EXIT_FAILURE;
                goto drain;
            }

            CloseHandle(WriteOutPipe);
            CloseHandle(WriteErrPipe);

            CurrentProcess++;
        }
    }

    //
    //  If we didn't find any source file, just execute the
    //  command verbatim.  Since there's only one child here we just
    //  let it do IO to this process output and error handles.
    //

    if (CurrentProcess == 0) {
        STARTUPINFO StartupInfo;
        DWORD MyProcess = CurrentProcess;

        ZeroMemory( &StartupInfo, sizeof(StartupInfo) );
        StartupInfo.cb = sizeof(StartupInfo);

        if (!CreateProcess( NULL, szCmdCommon, NULL, NULL, TRUE, 0, NULL, NULL, &StartupInfo, &ProcessInfo[MyProcess].WindowsProcessInfo )) {
            return EXIT_FAILURE;
        }

        CurrentProcess++;
    }

drain:

    //
    //  Now wait on all of the child processes.  Note that if any
    //  fail we will also fail with the same error code.
    //

    if (CurrentProcess > NumberProcesses) {
        CurrentProcess = NumberProcesses;
    }

    do {
        CurrentProcess--;

        WaitOnProcess( &ProcessInfo[CurrentProcess] );

    } while (CurrentProcess > 0);

    if (GlobalExitCode) {
        ExitProcess(GlobalExitCode);
    }

    //
    //  All of the child processes succeeded, so we did too.
    //

    return EXIT_SUCCESS;
}

// vim:sw=4:ts=4:et:
